#encoding=utf8
import numpy as np
import pickle

#********* Begin *********#
# 此处可选填

#********* End *********#

def load_dataset(file_name):
    '''
    从文件读入数据集
    被多处调用，请勿删除或改动本函数！！！
    '''
    with open(file_name, 'rb') as f:
        raw_dataset = pickle.load(f)
        example_image = raw_dataset[0][0]
    num_classes = len(raw_dataset)
    dataset = np.empty((0, example_image.size))
    labels = np.empty((0, 1))
    for i_class in range(num_classes):
        num_images = len(raw_dataset[i_class])
        for i_image in range(num_images):        
            image = raw_dataset[i_class][i_image]
            features = image.flatten()
            dataset = np.vstack((dataset, features))
            labels = np.vstack((labels, i_class))
    return dataset, labels

class Classifier:
    def __init__(self):  # 初始化模型和数据 
        self.model = None

        # 下行代码从training_dataset.pkl文件读入训练数据，得到：
        #（1）训练数据 train_dataset：形状为(1000, 784)的ndarray
        #（2）训练数据的标签 train_labels：形状为(1000, 1)的ndarray
        self.train_dataset, self.train_labels = load_dataset('./step1/training_dataset.pkl')

    def train(self): # 训练模型
        
        #********* BEGIN *********#
        # 此处必填。以下的逻辑回归代码仅供参考，需替换成你自己的代码
        from sklearn.linear_model import LogisticRegression
        self.model = LogisticRegression(C=0.99, solver='lbfgs', multi_class='multinomial', max_iter=2000)
        train_labels = self.train_labels.ravel()
        self.model.fit(self.train_dataset, train_labels)  # 训练

        #********* END *********#

    def predict(self, test_dataset):  # 预测
        '''
        输入：测试数据 test_dataset: 形状为(500, 784)的ndarray
        输出：预测结果 predicted_labels: 形状为(500, )的ndarray
        '''

        #********* BEGIN *********#
        # 此处必填。以下代码仅供参考，需替换成你自己的代码
        predicted_labels = self.model.predict(test_dataset)  # 预测

        #********* END *********#
        return predicted_labels
    
    #********* Begin *********#
    # 此处可选填

    #********* End *********#
    

#********* Begin *********#
# 此处可选填

#********* End *********#